# -*- coding: utf-8 -*-
# !/usr/bin/env python
"""
-------------------------------------------------
   File Name：     datalist
   Description :   
   Author :       lth
   date：          2023/1/30
-------------------------------------------------
   Change Activity:
                   2023/1/30 10:28: create this script
-------------------------------------------------
"""
__author__ = 'lth'

import numpy as np
from torch.utils.data import Dataset
from transformers import BertTokenizer

tokenizer = BertTokenizer.from_pretrained("bert-base-cased")
labels = {
    "business": 0,
    "entertainment": 1,
    "sport": 2,
    "tech": 3,
    "politics": 4
}


class TextClassifyDataset(Dataset):
    def __init__(self, df):
        super(TextClassifyDataset, self).__init__()
        self.labels = [labels[label] for label in df["category"]]
        self.texts = [tokenizer(text,
                                padding='max_length',
                                max_length=512,
                                truncation=True,
                                return_tensors="pt")
                      for text in df['text']]

    def classes(self):
        return self.labels

    def __len__(self):
        return len(self.labels)

    def __get_batch_labels(self, idx):
        return np.array(self.labels[idx])

    def __get_batch_texts(self, idx):
        return self.texts[idx]

    def __getitem__(self, index):
        batch_texts = self.__get_batch_texts(index)
        batch_labels = self.__get_batch_labels(index)
        return batch_texts, batch_labels
